package com.example.scheduler;

import com.example.mq.RocketMqProducer;
import com.example.mq.RocketMqConsumerService;
import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyContext;
import org.apache.rocketmq.client.consumer.listener.ConsumeConcurrentlyStatus;
import org.apache.rocketmq.common.message.MessageExt;

import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * TaskExecutor：执行器
 * - 可直接接受SchedulerNode的RPC调用 directExecute()
 * - 也可从MQ中批量拉取消息 handleTaskMessage()
 * - 使用快/慢线程池区别处理
 * - 完成后可回调调度器，以更新DAG/普通任务状态
 */
public class TaskExecutor {

    private final String executorId;
    private final RocketMqProducer mqProducer;
    private final RocketMqConsumerService consumerSvc;
    private final SchedulerNode scheduler;

    // 快线程池 / 慢线程池，用于隔离耗时较长任务
    private final ExecutorService fastPool;
    private final ExecutorService slowPool;

    private DefaultMQPushConsumer consumer;
    private final AtomicBoolean running = new AtomicBoolean(true);

    // MQ相关配置
    private final String consumerGroup;
    private final String topic = "TASK_TOPIC";

    public TaskExecutor(String executorId,
                        RocketMqProducer prod,
                        RocketMqConsumerService svc,
                        SchedulerNode scheduler) {
        this.executorId = executorId;
        this.mqProducer = prod;
        this.consumerSvc = svc;
        this.scheduler = scheduler;
        this.fastPool = Executors.newFixedThreadPool(2);
        this.slowPool = Executors.newFixedThreadPool(2);

        // 对应的ConsumerGroup
        this.consumerGroup = "batch_consumer_group_" + executorId;
    }

    public String getExecutorId() {
        return executorId;
    }

    /**
     * 启动MQ消费者，从TASK_TOPIC中批量拉取任务消息
     */
    public void startWorking() {
        try {
            consumer = consumerSvc.startConsumer(
                    consumerGroup, topic,
                    (List<MessageExt> msgs, ConsumeConcurrentlyContext context) -> {
                        System.out.println("[Executor " + executorId + "] Batch pull count: " + msgs.size());
                        for (MessageExt msg : msgs) {
                            handleTaskMessage(msg);
                        }
                        return ConsumeConcurrentlyStatus.CONSUME_SUCCESS;
                    }
            );
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 解析消息：例如 "TASK|SLOW|jobId123" 或 "TASK|FAST|jobIdABC"
     */
    private void handleTaskMessage(MessageExt msg) {
        if (!running.get()) return;
        try {
            String body = new String(msg.getBody(), "UTF-8");
            if (body.startsWith("TASK|")) {
                String[] parts = body.split("\\|");
                if (parts.length >= 3) {
                    boolean isSlow = "SLOW".equals(parts[1]);
                    String jobId = parts[2];
                    dispatchJobInternally(jobId, isSlow);
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * directExecute：普通模式下由调度器直接调用
     * @param jobId  任务ID(若是DAG子任务则形如 "dagId:subTaskId")
     * @param isSlow 是否慢任务
     */
    public void directExecute(String jobId, boolean isSlow) {
        System.out.println("[Executor " + executorId + "] directExecute => job: " + jobId
                + ", isSlow=" + isSlow);
        dispatchJobInternally(jobId, isSlow);
    }

    /**
     * 根据isSlow选择快/慢线程池
     */
    private void dispatchJobInternally(String jobId, boolean isSlow) {
        if (!running.get()) return;
        if (isSlow) {
            slowPool.execute(() -> handleTask(jobId));
        } else {
            fastPool.execute(() -> handleTask(jobId));
        }
    }

    /**
     * 实际执行任务（模拟随机耗时），执行完后回调Scheduler
     */
    private void handleTask(String jobId) {
        long start = System.currentTimeMillis();
        System.out.println("[Executor " + executorId + "] Start job: " + jobId);
        try {
            Thread.sleep(ThreadLocalRandom.current().nextInt(50, 200)); // 模拟执行耗时
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
        long cost = System.currentTimeMillis() - start;
        boolean success = (ThreadLocalRandom.current().nextInt(100) >= 15); // 85%成功率
        String result = success ? "OK" : "FAIL";
        System.out.println("[Executor " + executorId + "] Finish job: " + jobId
                + ", cost=" + cost + "ms, result=" + result);

        // 如果是DAG子任务，调用Scheduler的 processDagResult()
        if (jobId.contains(":") && scheduler != null) {
            scheduler.processDagResult(jobId, cost, result);
        } else {
            // 普通任务：可发 “RESULT|OK|jobId|COST=xxx” 消息到 RESULT_TOPIC
            // 这里简化不写
        }
    }

    public void shutdown() {
        running.set(false);
        if (consumer != null) {
            consumer.shutdown();
        }
        fastPool.shutdownNow();
        slowPool.shutdownNow();
        System.out.println("Executor " + executorId + " shutdown.");
    }
}